package com.ruoyi.ai.service.impl;

import com.ruoyi.ai.service.ChatService;
import com.ruoyi.ai.tools.ChatTools;
import com.ruoyi.ai.service.VectorService;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.tool.method.MethodToolCallback;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.ReflectionUtils;
import reactor.core.publisher.Flux;

import java.lang.reflect.Method;
import java.util.List;
import java.util.Map;

/**
 * @author coach tam
 * @email 327395128@qq.com
 * @values 坚持灵活 灵活坚持
 * @since 2025/3/22
 *
 *
 *          经过了三次修改：
 *          版本一：只有一些基础能力，没有记忆功能（无历史对话）
 *          版本二：有记忆功能，可以查询历史对话
 *          版本三：增加RAG检索增强功能，将知识库文档和AI模型结合
 *
 */
@Slf4j
@Service
public class ChatServiceImpl implements ChatService {

    public static final String PROMPT = """
            你是一个雅橙科技教学助手，由龙哥研制，正在回答用户的提问。
            所有的回答都只能通过中文回答。
            所有回答的格式都是纯文本格式。
            我们公司的口号是"用心做教育"，你只能回答编程问题以及教学相关的业务，。
            使用来自“文档”部分的信息来提供准确的答案,尽可能保留原“文档”的内容。如果不确定或在“文档”部分找不到答案，只需表明你不知道答案即可。
            历史对话：
            {history}
            
            文档:
            {documents}
            
            问题
            {input}
            
            """;
    @Autowired
    private OpenAiChatModel chatModel;


    @Autowired
    private ChatTools chatTools;

    // 配置 ChatMemory
    private ChatMemory chatMemory = new InMemoryChatMemory();
    // 构建聊天选项
    private ChatOptions chatOptions;

    @Autowired
    private VectorService vectorService;

    @PostConstruct
    public void init(){
        // 创建工具回调集合
        ToolCallback[] toolCallbacksChains = {
                createToolCallback(chatTools, "getCurrentDateTime", "获取用户当前的日期信息"),
                createToolCallback(chatTools, "queryTeachingPlan", "查询课表"),
                createToolCallback(chatTools, "route", "根据页面名称跳转到指定页面"),
                createToolCallback(chatTools, "queryAssistantName", "查询助手信息"),
                createToolCallback(chatTools, "queryUserName", "查询用户信息和学习情况"),
                createToolCallback(chatTools, "queryCourse", "查询课程信息"),
        };

        // 使用 ToolCallingChatOptions 构建器创建聊天选项对象
        chatOptions = ToolCallingChatOptions.builder()
                // 设置工具回调链
                .toolCallbacks(toolCallbacksChains)
                // 禁用内部工具执行功能
                .internalToolExecutionEnabled(false)
                // 设置生成文本的最大令牌数为 200
                .maxTokens(200)
                // 设置频率惩罚系数为 1.5，以控制重复内容的生成
                .frequencyPenalty(1.5)
                // 设置温度参数为 0.75，以平衡生成内容的随机性和确定性
                .temperature(0.75)
                .build();
    }

    @Override
    public Flux<ChatResponse> chatWithTools(String message) {
        return Flux.deferContextual(ctx -> {
            ToolCallingManager toolCallingManager = ToolCallingManager.builder().build();

            // 创建提示模板
            PromptTemplate promptTemplate = new PromptTemplate(PROMPT);

            String userId = ctx.get("userId");
            // 获取历史对话
            List<Message> historyMessages = chatMemory.get(userId, 10);

            // 获取当本地知识库文档
            List<String> contentList = vectorService.searchQuestion(message);
            // 构建提示并发送请求
            Map<String, Object> promptParams = Map.of(
                    "input", message,
                    "history", historyMessages,
                    "documents", contentList
            );
            Prompt prompt = promptTemplate.create(promptParams, chatOptions);

            ChatResponse chatResponse = chatModel.call(prompt);

            // 处理工具调用
            while (chatResponse.hasToolCalls()) {
                ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, chatResponse);
                // 如果是route方法的调用结果
                Message messageObj = toolExecutionResult.conversationHistory().get(1);
                if (messageObj instanceof AssistantMessage) {
                    AssistantMessage assistantMessage = (AssistantMessage) messageObj;
                    if (!assistantMessage.getToolCalls().isEmpty()) {
                        AssistantMessage.ToolCall toolCall = assistantMessage.getToolCalls().get(0);
                        if(toolCall.name().equals("route")){
                            // 从toolExecutionResult的对话历史中获取第三个对象
                            Message thirdMessage = toolExecutionResult.conversationHistory().get(2);
                            if (thirdMessage instanceof ToolResponseMessage) {
                                ToolResponseMessage toolResponseMessage = (ToolResponseMessage) thirdMessage;
                                List<ToolResponseMessage.ToolResponse> responses = toolResponseMessage.getResponses();
                                if (!responses.isEmpty()) {
                                    ToolResponseMessage.ToolResponse toolResponse = responses.get(0);
                                    String textContent = toolResponse.responseData();
                                    // 提取[]中的字符串
                                    String extractedText = "";
                                    int startIndex = textContent.indexOf("[");
                                    int endIndex = textContent.indexOf("]");
                                    if (startIndex != -1 && endIndex != -1 && startIndex < endIndex) {
                                        extractedText = textContent.substring(startIndex + 1, endIndex);
                                        if (!extractedText.startsWith("/")) {
                                            extractedText = "/" + extractedText;
                                        }
                                    }
                                    // 返回跳转指令
                                    Generation generation = new Generation(new AssistantMessage("跳转至：" + extractedText));
                                    return Flux.just(new ChatResponse(List.of(generation)));
                                }
                            }
                        }
                    }
                }
                // 先将 conversationHistory 转为列表
                List<Message> conversationHistory = toolExecutionResult.conversationHistory();
                // 创建一个新的 UserMessage 包含用户 ID
                UserMessage userIdMessage = new UserMessage(
                        "用户 ID为: " + userId +
                                "不用将用户id显示出来，就直接显示当前用户即可" +
                                "然后在回答用户时，每个大的分类前面加上对应的提示如：目前面试题掌握情况(0-10分)、每日任务完成情况、每周测试成绩" +
                                "最后一个要求：把格式调整一下，更加美观");
                // 将包含用户 ID 的消息添加到对话历史中
                conversationHistory.add(userIdMessage);
                // 创建新的 Prompt
                prompt = new Prompt(conversationHistory, chatOptions);

                chatResponse = chatModel.call(prompt);
            }

            // 流式输出结果
            Flux<ChatResponse> stream = chatModel.stream(prompt);

            // 拼接完整的结果
            StringBuilder sb = new StringBuilder();
            return stream.doOnNext(cr -> {
                //拼接完整的结果
                String text = cr.getResults().get(0).getOutput().getText();
                log.debug("接收到一条数据:{}", text);
                sb.append(text);
            }).doOnComplete(() -> {
                // 更新对话历史

                log.debug("接收完成,更新聊天記錄");
                chatMemory.add(userId, new UserMessage(message));
                chatMemory.add(userId, new AssistantMessage(sb.toString()));
            }).doFinally(signalType -> {
                // 无论成功还是失败都会执行
                log.debug("流处理结束，信号类型:{}", signalType);
            });
        });
    }

    /**
     * 创建工具回调
     *
     * @param toolObject  工具对象
     * @param methodName  方法名
     * @param description 描述
     * @return 工具回调
     */
    private ToolCallback createToolCallback(Object toolObject, String methodName, String description) {
        // 先查找带String参数的方法
        Method method = ReflectionUtils.findMethod(toolObject.getClass(), methodName, String.class);
        if (method == null) {
            // 如果没找到带String参数的方法，查找无参数的方法
            method = ReflectionUtils.findMethod(toolObject.getClass(), methodName);
            if (method == null) {
                throw new IllegalArgumentException("Method not found: " + methodName + "(String) or " + methodName + "()");
            }
        }
        return MethodToolCallback.builder()
                .toolDefinition(ToolDefinition.builder(method).description(description).build())
                .toolMethod(method)
                .toolObject(toolObject)
                .build();
    }
}
